{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "98188280",
   "metadata": {},
   "source": [
    "## 原文与译文词表"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "1443aa2d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "vocab_x= {'<SOS>': 0, '<EOS>': 1, '<PAD>': 2, '0': 3, '1': 4, '2': 5, '3': 6, '4': 7, '5': 8, '6': 9, '7': 10, '8': 11, '9': 12, 'q': 13, 'w': 14, 'e': 15, 'r': 16, 't': 17, 'y': 18, 'u': 19, 'i': 20, 'o': 21, 'p': 22, 'a': 23, 's': 24, 'd': 25, 'f': 26, 'g': 27, 'h': 28, 'j': 29, 'k': 30, 'l': 31, 'z': 32, 'x': 33, 'c': 34, 'v': 35, 'b': 36, 'n': 37, 'm': 38}\n",
      "vocab_y= {'<SOS>': 0, '<EOS>': 1, '<PAD>': 2, '0': 3, '1': 4, '2': 5, '3': 6, '4': 7, '5': 8, '6': 9, '7': 10, '8': 11, '9': 12, 'Q': 13, 'W': 14, 'E': 15, 'R': 16, 'T': 17, 'Y': 18, 'U': 19, 'I': 20, 'O': 21, 'P': 22, 'A': 23, 'S': 24, 'D': 25, 'F': 26, 'G': 27, 'H': 28, 'J': 29, 'K': 30, 'L': 31, 'Z': 32, 'X': 33, 'C': 34, 'V': 35, 'B': 36, 'N': 37, 'M': 38}\n"
     ]
    }
   ],
   "source": [
    "# 第13章/定义字典\n",
    "vocab_x = \"<SOS>,<EOS>,<PAD>,0,1,2,3,4,5,6,7,8,9,q,w,e,r,t,y,u,i,o,p,a,s,d,f,g,h,j,k,l,z,x,c,v,b,n,m\"\n",
    "\n",
    "# 原文：单词和索引的对应关系\n",
    "vocab_x = {word: i for i, word in enumerate(vocab_x.split(\",\"))}\n",
    "\n",
    "# 原文单词表\n",
    "vocab_xr = [k for k, v in vocab_x.items()]\n",
    "\n",
    "# 译文：单词和索引的对应关系\n",
    "vocab_y = {k.upper(): v for k, v in vocab_x.items()}\n",
    "\n",
    "# 译文单词表\n",
    "vocab_yr = [k for k, v in vocab_y.items()]\n",
    "\n",
    "print(\"vocab_x=\", vocab_x)\n",
    "print(\"vocab_y=\", vocab_y)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4fdfae6",
   "metadata": {},
   "source": [
    "## 生成一个样本"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "cee56575",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([ 0,  5, 36, 25, 26, 25, 27, 36, 33, 29, 10, 35, 28, 17, 30, 32, 30, 27,\n",
       "         32, 22, 32, 32, 22, 36, 23, 10, 12, 34, 34, 25, 28, 30, 26, 29, 37, 27,\n",
       "          1,  2,  2,  2,  2,  2,  2,  2,  2,  2,  2,  2,  2,  2]),\n",
       " tensor([ 0, 27, 27, 37, 29, 26, 30, 28, 25, 34, 34,  3,  5, 23, 36, 22, 32, 32,\n",
       "         22, 32, 27, 30, 32, 30, 17, 28, 35,  5, 29, 33, 36, 27, 25, 26, 25, 36,\n",
       "         10,  1,  2,  2,  2,  2,  2,  2,  2,  2,  2,  2,  2,  2,  2]))"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#第13章/定义生成数据函数\n",
    "import random\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "\n",
    "\n",
    "def get_data():\n",
    "    \"\"\"生成一条示例, x表示原文, y表示译文\"\"\"\n",
    "\n",
    "    # 定义 原文词集合\n",
    "    words = [\n",
    "        '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'q', 'w', 'e', 'r',\n",
    "        't', 'y', 'u', 'i', 'o', 'p', 'a', 's', 'd', 'f', 'g', 'h', 'j', 'k',\n",
    "        'l', 'z', 'x', 'c', 'v', 'b', 'n', 'm'\n",
    "    ]\n",
    "\n",
    "    # 定义每个词被选中的概率\n",
    "    p = np.array([\n",
    "        1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12,\n",
    "        13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26\n",
    "    ])\n",
    "    p = p / p.sum()\n",
    "\n",
    "    # 随机选n个词\n",
    "    n = random.randint(30, 48)\n",
    "    x = np.random.choice(words, size=n, replace=True, p=p)\n",
    "\n",
    "    # 采样的结果就是x\n",
    "    x = x.tolist()\n",
    "\n",
    "    # y是对x的变换得到的\n",
    "    # 字母大写,数字取9以内的互补数\n",
    "    def f(i):\n",
    "        i = i.upper()\n",
    "        if not i.isdigit():\n",
    "            return i\n",
    "        i = 9 - int(i)\n",
    "        return str(i)\n",
    "\n",
    "    y = [f(i) for i in x]\n",
    "    # 逆序\n",
    "    y = y[::-1]\n",
    "\n",
    "    # y中的首字母双写\n",
    "    y = [y[0]] + y\n",
    "\n",
    "    # 加上首尾符号\n",
    "    x = ['<SOS>'] + x + ['<EOS>']\n",
    "    y = ['<SOS>'] + y + ['<EOS>']\n",
    "\n",
    "    # 补pad到固定长度\n",
    "    x = x + ['<PAD>'] * 50\n",
    "    y = y + ['<PAD>'] * 51\n",
    "    x = x[:50]\n",
    "    y = y[:51]\n",
    "\n",
    "    # 编码成数据\n",
    "    x = [vocab_x[i] for i in x]\n",
    "    y = [vocab_y[i] for i in y]\n",
    "\n",
    "    # 转tensor\n",
    "    x = torch.LongTensor(x)\n",
    "    y = torch.LongTensor(y)\n",
    "\n",
    "    return x, y\n",
    "\n",
    "get_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "b239c926",
   "metadata": {},
   "outputs": [],
   "source": [
    "# #第13章/两数相加测试\n",
    "# #使用这份数据请把训练次数改为10\n",
    "# def get_data():\n",
    "#     #定义词集合\n",
    "#     words = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9']\n",
    "\n",
    "#     #定义每个词被选中的概率\n",
    "#     p = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10])\n",
    "#     p = p / p.sum()\n",
    "\n",
    "#     #随机选n个词\n",
    "#     n = random.randint(10, 20)\n",
    "#     s1 = np.random.choice(words, size=n, replace=True, p=p)\n",
    "\n",
    "#     #采样的结果就是s1\n",
    "#     s1 = s1.tolist()\n",
    "\n",
    "#     #同样的方法,再采出s2\n",
    "#     n = random.randint(10, 20)\n",
    "#     s2 = np.random.choice(words, size=n, replace=True, p=p)\n",
    "#     s2 = s2.tolist()\n",
    "\n",
    "#     #y等于s1和s2数值上的相加\n",
    "#     y = int(''.join(s1)) + int(''.join(s2))\n",
    "#     y = list(str(y))\n",
    "\n",
    "#     #x等于s1和s2字符上的相加\n",
    "#     x = s1 + ['a'] + s2\n",
    "\n",
    "#     #加上首尾符号\n",
    "#     x = ['<SOS>'] + x + ['<EOS>']\n",
    "#     y = ['<SOS>'] + y + ['<EOS>']\n",
    "\n",
    "#     #补pad到固定长度\n",
    "#     x = x + ['<PAD>'] * 50\n",
    "#     y = y + ['<PAD>'] * 51\n",
    "#     x = x[:50]\n",
    "#     y = y[:51]\n",
    "\n",
    "#     #编码成数据\n",
    "#     x = [vocab_x[i] for i in x]\n",
    "#     y = [vocab_y[i] for i in y]\n",
    "\n",
    "#     #转tensor\n",
    "#     x = torch.LongTensor(x)\n",
    "#     y = torch.LongTensor(y)\n",
    "\n",
    "#     return x, y\n",
    "\n",
    "\n",
    "# get_data()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "850a349e",
   "metadata": {},
   "source": [
    "## 定义数据集和加载器"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "99e7c73d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([8, 50]), torch.Size([8, 51]))"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 第13章/定义数据集和加载器\n",
    "# 定义数据集\n",
    "class Dataset(torch.utils.data.Dataset):\n",
    "    def __init__(self):\n",
    "\n",
    "        # 父类初始化方法\n",
    "        super(Dataset, self).__init__()\n",
    "\n",
    "    def __len__(self):\n",
    "\n",
    "        # 生成 10000 条数据\n",
    "        return 10000\n",
    "\n",
    "    def __getitem__(self, i):\n",
    "\n",
    "        # 返回 一个样本\n",
    "        return get_data()\n",
    "\n",
    "\n",
    "# 数据加载器\n",
    "loader = torch.utils.data.DataLoader(\n",
    "    dataset=Dataset(),  # 数据集\n",
    "    batch_size=8,\n",
    "    drop_last=True,  # 如果在最后一个批次中数据量不足 batch_size，则丢弃该批次\n",
    "    shuffle=True,  # 打乱数据\n",
    "    collate_fn=None,  # 用于处理一批数据样本的自定义函数\n",
    ")\n",
    "\n",
    "# 查看数据样例\n",
    "for i, (x, y) in enumerate(loader):\n",
    "    break\n",
    "\n",
    "x.shape, y.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "85e33391",
   "metadata": {},
   "source": [
    "## 定义 PAD MASK"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "1ae1b754",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[[False, False, False,  ...,  True,  True,  True],\n",
       "          [False, False, False,  ...,  True,  True,  True],\n",
       "          [False, False, False,  ...,  True,  True,  True],\n",
       "          ...,\n",
       "          [False, False, False,  ...,  True,  True,  True],\n",
       "          [False, False, False,  ...,  True,  True,  True],\n",
       "          [False, False, False,  ...,  True,  True,  True]]]])"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#第13章/定义mask_pad函数\n",
    "def mask_pad(data):\n",
    "    # 一批次训练数据，有 batch_size 句话, 每句话50个词, 这里是还没embed的\n",
    "    # data = [batch_size, 50]\n",
    "\n",
    "    # 判断每个词是不是 <PAD>\n",
    "    mask = data == vocab_x['<PAD>']\n",
    "\n",
    "    #[batch_size, 50] -> [batch_size, 1, 1, 50]\n",
    "    mask = mask.reshape(-1, 1, 1, 50)\n",
    "\n",
    "    # 在计算一个样本的注意力时, 是计算每一个词和其余词之间的注意力, 所以是个 50*50 的矩阵\n",
    "    # 是 pad的列 是 true, 意味着任何词对pad的注意力都是 0\n",
    "    # 但是 pad 本身对其他词的注意力并不是0, 所以是 pad 的行不是 true\n",
    "\n",
    "    # 复制n次\n",
    "    # [b, 1, 1, 50] -> [b, 1, 50, 50]\n",
    "    mask = mask.expand(-1, 1, 50, 50)\n",
    "\n",
    "    return mask\n",
    "\n",
    "\n",
    "mask_pad(x[:1])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a792bbd",
   "metadata": {},
   "source": [
    "## 上三角 MASK"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "dc80f8eb",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[[False,  True,  True,  ...,  True,  True,  True],\n",
       "          [False, False,  True,  ...,  True,  True,  True],\n",
       "          [False, False, False,  ...,  True,  True,  True],\n",
       "          ...,\n",
       "          [False, False, False,  ...,  True,  True,  True],\n",
       "          [False, False, False,  ...,  True,  True,  True],\n",
       "          [False, False, False,  ...,  True,  True,  True]]]])"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#第13章/定义mask_tril函数\n",
    "def mask_tril(data):\n",
    "    # batch_size句话, 每句话50个词, 这里是还没embed的\n",
    "    # data = [batch_size, 50]\n",
    "\n",
    "    # 50*50的矩阵表示每个词对其他词是否可见\n",
    "    # 上三角矩阵,不包括对角线,意味着,对每个词而言,他只能看到他自己,和他之前的词,而看不到之后的词\n",
    "    # [1, 50, 50]\n",
    "    \"\"\"\n",
    "    [[0, 1, 1, 1, 1],\n",
    "     [0, 0, 1, 1, 1],\n",
    "     [0, 0, 0, 1, 1],\n",
    "     [0, 0, 0, 0, 1],\n",
    "     [0, 0, 0, 0, 0]]\"\"\"\n",
    "    tril = 1 - torch.tril(torch.ones(1, 50, 50, dtype=torch.long))\n",
    "\n",
    "    # 判断y当中每个词是不是pad, 如果是pad则不可见\n",
    "    # [batch_size, 50]\n",
    "    mask = data == vocab_y['<PAD>']\n",
    "\n",
    "    # 变形+转型,为了之后的计算\n",
    "    # [b, 1, 50]\n",
    "    mask = mask.unsqueeze(1).long()\n",
    "\n",
    "    # mask和tril求并集\n",
    "    # [b, 1, 50] + [1, 50, 50] -> [b, 50, 50]\n",
    "    mask = mask + tril\n",
    "\n",
    "    # 转布尔型\n",
    "    mask = mask > 0\n",
    "\n",
    "    # 转布尔型,增加一个维度,便于后续的计算\n",
    "    mask = (mask == 1).unsqueeze(dim=1)\n",
    "\n",
    "    return mask\n",
    "\n",
    "mask_tril(x[:1])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9747fbd",
   "metadata": {},
   "source": [
    "## 多头注意力计算函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "a59b48af",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([8, 50, 32])"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#第13章/定义注意力计算函数\n",
    "def attention(Q, K, V, mask):\n",
    "    # batch_size 句话, 每句话50个词, 每个词编码成32维向量, 4个头,每个头分到8维向量\n",
    "    # Q、K、V = [b, 4, 50, 8]\n",
    "\n",
    "    # [b, 4, 50, 8] * [b, 4, 8, 50] -> [b, 4, 50, 50]\n",
    "    # Q、K矩阵相乘,求每个词相对其他所有词的注意力\n",
    "    score = torch.matmul(Q, K.permute(0, 1, 3, 2))\n",
    "\n",
    "    # 除以每个头维数的平方根,做数值缩放\n",
    "    score /= 8**0.5\n",
    "\n",
    "    # mask遮盖, mask是true的地方都被替换成-inf, 这样在计算softmax的时候, -inf会被压缩到0\n",
    "    # mask = [batch_size, 1, 50, 50]\n",
    "    score = score.masked_fill_(mask, -float('inf'))\n",
    "    score = torch.softmax(score, dim=-1)\n",
    "\n",
    "    # 以注意力分数乘以V, 得到最终的注意力结果\n",
    "    # [b, 4, 50, 50] * [b, 4, 50, 8] -> [b, 4, 50, 8]\n",
    "    score = torch.matmul(score, V)\n",
    "\n",
    "    # 每个头计算的结果合一\n",
    "    # [b, 4, 50, 8] -> [b, 50, 32]\n",
    "    score = score.permute(0, 2, 1, 3).reshape(-1, 50, 32)\n",
    "\n",
    "    return score\n",
    "\n",
    "\n",
    "attention(torch.randn(8, 4, 50, 8), torch.randn(8, 4, 50, 8),\n",
    "            torch.randn(8, 4, 50, 8), torch.zeros(8, 1, 50, 50)).shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ff8cae4f",
   "metadata": {},
   "source": [
    "## LayerNorm & BatchNorm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "d308d219",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[-1.1761, -1.0523, -0.9285, -0.8047],\n",
      "         [-1.1761, -1.0523, -0.9285, -0.8047],\n",
      "         [-1.1761, -1.0523, -0.9285, -0.8047],\n",
      "         [-1.1761, -1.0523, -0.9285, -0.8047]],\n",
      "\n",
      "        [[ 0.8047,  0.9285,  1.0523,  1.1761],\n",
      "         [ 0.8047,  0.9285,  1.0523,  1.1761],\n",
      "         [ 0.8047,  0.9285,  1.0523,  1.1761],\n",
      "         [ 0.8047,  0.9285,  1.0523,  1.1761]]],\n",
      "       grad_fn=<NativeBatchNormBackward0>)\n",
      "tensor([[[-1.3416, -0.4472,  0.4472,  1.3416],\n",
      "         [-1.3416, -0.4472,  0.4472,  1.3416],\n",
      "         [-1.3416, -0.4472,  0.4472,  1.3416],\n",
      "         [-1.3416, -0.4472,  0.4472,  1.3416]],\n",
      "\n",
      "        [[-1.3416, -0.4472,  0.4472,  1.3416],\n",
      "         [-1.3416, -0.4472,  0.4472,  1.3416],\n",
      "         [-1.3416, -0.4472,  0.4472,  1.3416],\n",
      "         [-1.3416, -0.4472,  0.4472,  1.3416]]],\n",
      "       grad_fn=<NativeLayerNormBackward0>)\n"
     ]
    }
   ],
   "source": [
    "# 第13章/BatchNorm1d和LayerNorm的对比\n",
    "# 标准化之后,均值是0,标准差是1\n",
    "# BN是取不同样本做标准化\n",
    "# LN是取不同通道做标准化\n",
    "# affine=True, elementwise_affine=True, 指定标准化后, 再计算一个线性映射\n",
    "norm = torch.nn.BatchNorm1d(num_features=4, affine=True)\n",
    "print(norm(torch.arange(32, dtype=torch.float32).reshape(2, 4, 4)))\n",
    "\n",
    "norm = torch.nn.LayerNorm(normalized_shape=4, elementwise_affine=True)\n",
    "print(norm(torch.arange(32, dtype=torch.float32).reshape(2, 4, 4)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "deb30419",
   "metadata": {},
   "source": [
    "## 多头注意力计算层"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "9a6b55f0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([8, 50, 32])"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#第13章/多头注意力计算层\n",
    "class MultiHead(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "\n",
    "        # 父类初始化\n",
    "        super().__init__()\n",
    "\n",
    "        # 全连接层\n",
    "        self.fc_Q = torch.nn.Linear(32, 32)\n",
    "        self.fc_K = torch.nn.Linear(32, 32)\n",
    "        self.fc_V = torch.nn.Linear(32, 32)\n",
    "\n",
    "        self.out_fc = torch.nn.Linear(32, 32)\n",
    "\n",
    "        # 标准化层\n",
    "        self.norm = torch.nn.LayerNorm(normalized_shape=32,\n",
    "                                    elementwise_affine=True)\n",
    "\n",
    "        # Dropout层\n",
    "        self.dropout = torch.nn.Dropout(p=0.1)\n",
    "\n",
    "    def forward(self, Q, K, V, mask):\n",
    "        \"\"\"前向传播\"\"\"\n",
    "\n",
    "        # batch_size句话, 每句话50个词, 每个词编码成32维向量\n",
    "        # Q、K、V = [b, 50, 32]\n",
    "        b = Q.shape[0]\n",
    "\n",
    "        # 保留下原始的Q, 后面要做短接用\n",
    "        clone_Q = Q.clone()\n",
    "\n",
    "        # 标准化\n",
    "        Q = self.norm(Q)\n",
    "        K = self.norm(K)\n",
    "        V = self.norm(V)\n",
    "\n",
    "        # 线性运算,维度不变\n",
    "        # [b, 50, 32] -> [b, 50, 32]\n",
    "        K = self.fc_K(K)\n",
    "        V = self.fc_V(V)\n",
    "        Q = self.fc_Q(Q)\n",
    "\n",
    "        # 拆分成多个头\n",
    "        # b句话,每句话50个词,每个词编码成32维向量,4个头,每个头分到8维向量\n",
    "        # [b, 50, 32] -> [b, 4, 50, 8]\n",
    "        Q = Q.reshape(b, 50, 4, 8).permute(0, 2, 1, 3)\n",
    "        K = K.reshape(b, 50, 4, 8).permute(0, 2, 1, 3)\n",
    "        V = V.reshape(b, 50, 4, 8).permute(0, 2, 1, 3)\n",
    "\n",
    "        # 计算注意力\n",
    "        # [b, 4, 50, 8] -> [b, 50, 32]\n",
    "        score = attention(Q, K, V, mask)\n",
    "\n",
    "        # 计算输出,维度不变\n",
    "        # [b, 50, 32] -> [b, 50, 32]\n",
    "        score = self.dropout(self.out_fc(score))\n",
    "\n",
    "        # 短接\n",
    "        score = clone_Q + score\n",
    "        return score\n",
    "\n",
    "\n",
    "MultiHead()(torch.randn(8, 50, 32), torch.randn(8, 50, 32),\n",
    "            torch.randn(8, 50, 32), torch.zeros(8, 1, 50, 50)).shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e2008c8",
   "metadata": {},
   "source": [
    "## 定义位置编码层"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "ade4476e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([8, 50, 32])"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#第13章/定义位置编码层\n",
    "import math\n",
    "\n",
    "\n",
    "class PositionEmbedding(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "\n",
    "        # 父类初始化方法\n",
    "        super().__init__()\n",
    "\n",
    "        # pos是第几个词, i是第几个维度, d_model是维度总数\n",
    "        def get_pe(pos, i, d_model):\n",
    "            \"\"\"获取位置编码矩阵\"\"\"\n",
    "            d = 1e4**(i / d_model)\n",
    "            pe = pos / d\n",
    "\n",
    "            if i % 2 == 0:\n",
    "                return math.sin(pe)\n",
    "            return math.cos(pe)\n",
    "\n",
    "        # 初始化位置编码矩阵\n",
    "        pe = torch.empty(50, 32)\n",
    "        for i in range(50):\n",
    "            for j in range(32):\n",
    "                pe[i, j] = get_pe(i, j, 32)\n",
    "        pe = pe.unsqueeze(0)\n",
    "\n",
    "        # 定义为不更新的常量\n",
    "        self.register_buffer('pe', pe)\n",
    "\n",
    "        # 词编码层\n",
    "        self.embed = torch.nn.Embedding(39, 32)\n",
    "        \n",
    "        # 初始化参数\n",
    "        self.embed.weight.data.normal_(0, 0.1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        #[8, 50] -> [8, 50, 32]\n",
    "        embed = self.embed(x)\n",
    "\n",
    "        # 词编码和位置编码相加\n",
    "        # [8, 50, 32] + [1, 50, 32] -> [8, 50, 32]\n",
    "        embed = embed + self.pe\n",
    "        return embed\n",
    "\n",
    "\n",
    "PositionEmbedding()(torch.ones(8, 50).long()).shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0644bae9",
   "metadata": {},
   "source": [
    "## 全连接输出层"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "d5892664",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([8, 50, 32])"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 第13章/定义全连接输出层\n",
    "class FullyConnectedOutput(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "\n",
    "        # 父类初始化\n",
    "        super().__init__()\n",
    "\n",
    "        # 顺序模型\n",
    "        self.fc = torch.nn.Sequential(\n",
    "            torch.nn.Linear(in_features=32, out_features=64),\n",
    "            torch.nn.ReLU(),\n",
    "            torch.nn.Linear(in_features=64, out_features=32),\n",
    "            torch.nn.Dropout(p=0.1),\n",
    "        )\n",
    "\n",
    "        # 标准化\n",
    "        self.norm = torch.nn.LayerNorm(normalized_shape=32,\n",
    "                                       elementwise_affine=True)\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"前向传播\"\"\"\n",
    "        \n",
    "        # 保留下原始的x,后面要做短接用\n",
    "        clone_x = x.clone()\n",
    "\n",
    "        # 标准化\n",
    "        x = self.norm(x)\n",
    "\n",
    "        # 线性全连接运算\n",
    "        # [b, 50, 32] -> [b, 50, 32]\n",
    "        out = self.fc(x)\n",
    "\n",
    "        # 做短接\n",
    "        out = clone_x + out\n",
    "\n",
    "        return out\n",
    "\n",
    "\n",
    "FullyConnectedOutput()(torch.randn(8, 50, 32)).shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc5b8ae2",
   "metadata": {},
   "source": [
    "## 定义编码器"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "86d18572",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([8, 50, 32])"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#第13章/定义编码器\n",
    "# 编码器层\n",
    "class EncoderLayer(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "\n",
    "        # 父类初始化\n",
    "        super().__init__()\n",
    "\n",
    "        # 多头注意力层\n",
    "        self.mh = MultiHead()\n",
    "\n",
    "        # 全连接输出层\n",
    "        self.fc = FullyConnectedOutput()\n",
    "\n",
    "    def forward(self, x, mask):\n",
    "        \"\"\"前向传播\"\"\"\n",
    "        # 计算自注意力,维度不变\n",
    "        # [b, 50, 32] -> [b, 50, 32]\n",
    "        score = self.mh(x, x, x, mask)\n",
    "\n",
    "        # 全连接输出,维度不变\n",
    "        #[b, 50, 32] -> [b, 50, 32]\n",
    "        out = self.fc(score)\n",
    "\n",
    "        return out\n",
    "\n",
    "\n",
    "# 编码器 -- 由多个编码层组成\n",
    "class Encoder(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.layer_1 = EncoderLayer()\n",
    "        self.layer_2 = EncoderLayer()\n",
    "        self.layer_3 = EncoderLayer()\n",
    "\n",
    "    def forward(self, x, mask):\n",
    "        x = self.layer_1(x, mask)\n",
    "        x = self.layer_2(x, mask)\n",
    "        x = self.layer_3(x, mask)\n",
    "        return x\n",
    "\n",
    "\n",
    "Encoder()(torch.randn(8, 50, 32), torch.ones(8, 1, 50, 50)).shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eac58b2a",
   "metadata": {},
   "source": [
    "## 定义解码器"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "6553c927",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([8, 50, 32])"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#第13章/定义解码器\n",
    "# 解码器层\n",
    "class DecoderLayer(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "\n",
    "        # 父类初始化\n",
    "        super().__init__()\n",
    "\n",
    "        # 多头注意力计算层\n",
    "        self.mh1 = MultiHead()\n",
    "        self.mh2 = MultiHead()\n",
    "\n",
    "        # 全连接输出层\n",
    "        self.fc = FullyConnectedOutput()\n",
    "\n",
    "    def forward(self, x, y, mask_pad_x, mask_tril_y):\n",
    "        # 先计算y的自注意力,维度不变\n",
    "        # [b, 50, 32] -> [b, 50, 32]\n",
    "        y = self.mh1(y, y, y, mask_tril_y)\n",
    "\n",
    "        # 结合x和y的注意力计算,维度不变\n",
    "        # [b, 50, 32],[b, 50, 32] -> [b, 50, 32]\n",
    "        y = self.mh2(y, x, x, mask_pad_x)\n",
    "\n",
    "        # 全连接输出,维度不变\n",
    "        # [b, 50, 32] -> [b, 50, 32]\n",
    "        y = self.fc(y)\n",
    "\n",
    "        return y\n",
    "\n",
    "\n",
    "# 解码器 -- 由多个解码器组成\n",
    "class Decoder(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "\n",
    "        # 父类初始化方法\n",
    "        super().__init__()\n",
    "\n",
    "        # 多个解码器层\n",
    "        self.layer_1 = DecoderLayer()\n",
    "        self.layer_2 = DecoderLayer()\n",
    "        self.layer_3 = DecoderLayer()\n",
    "\n",
    "    def forward(self, x, y, mask_pad_x, mask_tril_y):\n",
    "        \"\"\"前向传播\"\"\"\n",
    "\n",
    "        y = self.layer_1(x, y, mask_pad_x, mask_tril_y)\n",
    "        y = self.layer_2(x, y, mask_pad_x, mask_tril_y)\n",
    "        y = self.layer_3(x, y, mask_pad_x, mask_tril_y)\n",
    "        return y\n",
    "\n",
    "\n",
    "Decoder()(torch.randn(8, 50, 32), torch.randn(8, 50, 32),\n",
    "            torch.ones(8, 1, 50, 50), torch.ones(8, 1, 50, 50)).shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "65ea4f82",
   "metadata": {},
   "source": [
    "## 定义 Transformer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "b51b9f91",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([8, 50, 39])"
      ]
     },
     "execution_count": 46,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#第13章/定义主模型\n",
    "class Transformer(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "\n",
    "        # 位置编码层\n",
    "        self.embed_x = PositionEmbedding()\n",
    "        self.embed_y = PositionEmbedding()\n",
    "\n",
    "        # 编码器\n",
    "        self.encoder = Encoder()\n",
    "\n",
    "        # 解码器\n",
    "        self.decoder = Decoder()\n",
    "\n",
    "        # 全连接输出层\n",
    "        self.fc_out = torch.nn.Linear(32, 39)\n",
    "\n",
    "    def forward(self, x, y):\n",
    "        \"\"\"前向传播\"\"\"\n",
    "\n",
    "        # [b, 1, 50, 50]\n",
    "        mask_pad_x = mask_pad(x)\n",
    "        mask_tril_y = mask_tril(y)\n",
    "\n",
    "        # 编码,添加位置信息\n",
    "        # x = [b, 50] -> [b, 50, 32]\n",
    "        # y = [b, 50] -> [b, 50, 32]\n",
    "        x, y = self.embed_x(x), self.embed_y(y)\n",
    "\n",
    "        # 编码层计算\n",
    "        # [b, 50, 32] -> [b, 50, 32]\n",
    "        x = self.encoder(x, mask_pad_x)\n",
    "\n",
    "        # 解码层计算\n",
    "        # [b, 50, 32],[b, 50, 32] -> [b, 50, 32]\n",
    "        y = self.decoder(x, y, mask_pad_x, mask_tril_y)\n",
    "\n",
    "        # 全连接输出,维度不变\n",
    "        # [b, 50, 32] -> [b, 50, 39]\n",
    "        y = self.fc_out(y)\n",
    "\n",
    "        return y\n",
    "\n",
    "\n",
    "model = Transformer()\n",
    "\n",
    "model(torch.ones(8, 50).long(), torch.ones(8, 50).long()).shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1087bdc7",
   "metadata": {},
   "source": [
    "## 定义预测函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "9ff7f830",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 0, 28, 28, 28, 28, 28, 28, 28, 28, 28, 28, 31, 31, 31, 37, 37, 37, 37,\n",
       "         37, 37, 37, 37, 37, 37, 37, 37, 37, 37, 37, 37, 37, 37, 37, 37, 37, 37,\n",
       "         37, 37, 37, 37, 37, 37,  9,  9, 37, 37, 37, 37, 37, 37]])"
      ]
     },
     "execution_count": 47,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#第13章/定义预测函数\n",
    "def predict(x):\n",
    "    # x = [1, 50]\n",
    "    model.eval()\n",
    "\n",
    "    # [1, 1, 50, 50]\n",
    "    mask_pad_x = mask_pad(x)\n",
    "\n",
    "    # 初始化输出,这个是固定值\n",
    "    # [1, 50]\n",
    "    # [[0,2,2,2...]]\n",
    "    target = [vocab_y['<SOS>']] + [vocab_y['<PAD>']] * 49\n",
    "    target = torch.LongTensor(target).unsqueeze(0)\n",
    "\n",
    "    # x编码,添加位置信息\n",
    "    # [1, 50] -> [1, 50, 32]\n",
    "    x = model.embed_x(x)\n",
    "\n",
    "    # 编码层计算,维度不变\n",
    "    # [1, 50, 32] -> [1, 50, 32]\n",
    "    x = model.encoder(x, mask_pad_x)\n",
    "\n",
    "    # 遍历生成第1个词到第49个词\n",
    "    for i in range(49):\n",
    "        #[1, 50]\n",
    "        y = target\n",
    "\n",
    "        #[1, 1, 50, 50]\n",
    "        mask_tril_y = mask_tril(y)\n",
    "\n",
    "        #y编码,添加位置信息\n",
    "        #[1, 50] -> [1, 50, 32]\n",
    "        y = model.embed_y(y)\n",
    "\n",
    "        #解码层计算,维度不变\n",
    "        #[1, 50, 32],[1, 50, 32] -> [1, 50, 32]\n",
    "        y = model.decoder(x, y, mask_pad_x, mask_tril_y)\n",
    "\n",
    "        #全连接输出,39分类\n",
    "        #[1, 50, 32] -> [1, 50, 39]\n",
    "        out = model.fc_out(y)\n",
    "\n",
    "        #取出当前词的输出\n",
    "        #[1, 50, 39] -> [1, 39]\n",
    "        out = out[:, i, :]\n",
    "\n",
    "        #取出分类结果\n",
    "        #[1, 39] -> [1]\n",
    "        out = out.argmax(dim=1).detach()\n",
    "\n",
    "        #以当前词预测下一个词,填到结果中\n",
    "        target[:, i + 1] = out\n",
    "\n",
    "    return target\n",
    "\n",
    "\n",
    "predict(torch.ones(1, 50).long())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b3444790",
   "metadata": {},
   "source": [
    "## 定义训练函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "01693ce1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 0 0.002 3.773591995239258 0.06845238095238096\n",
      "1 0 0.002 0.1717977523803711 0.9735099337748344\n",
      "2 0 0.002 0.017590820789337158 1.0\n",
      "3 0 0.001 0.0005295049049891531 1.0\n",
      "4 0 0.001 0.006369311828166246 1.0\n",
      "5 0 0.001 0.009850824251770973 1.0\n",
      "6 0 0.0005 0.0008660162566229701 1.0\n",
      "7 0 0.0005 0.06849044561386108 0.9793510324483776\n",
      "8 0 0.0005 0.00019880352192558348 1.0\n",
      "9 0 0.00025 0.0010993623873218894 1.0\n"
     ]
    }
   ],
   "source": [
    "#第13章/定义训练函数\n",
    "def train():\n",
    "\n",
    "    # 交叉熵损失\n",
    "    loss_func = torch.nn.CrossEntropyLoss()\n",
    "\n",
    "    # 优化器\n",
    "    optim = torch.optim.Adam(model.parameters(), lr=2e-3)\n",
    "\n",
    "    # 步长调整学习率\n",
    "    sched = torch.optim.lr_scheduler.StepLR(optim, # 优化器对象\n",
    "                                            step_size=3, # 调整周期\n",
    "                                            gamma=0.5) # 学习率调整幅度\n",
    "\n",
    "    # 10个 epoch\n",
    "    for epoch in range(10):\n",
    "\n",
    "        # 遍历数据集\n",
    "        for i, (x, y) in enumerate(loader):\n",
    "            # x = [8, 50]\n",
    "            # y = [8, 51]\n",
    "\n",
    "            # 在训练时, 是拿y的每一个字符输入, 预测下一个字符,所以不需要最后一个字\n",
    "            # [8, 50, 39]\n",
    "            pred = model(x, y[:, :-1])\n",
    "\n",
    "            # [8, 50, 39] -> [400, 39]\n",
    "            pred = pred.reshape(-1, 39)\n",
    "\n",
    "            # [8, 51] -> [400]\n",
    "            y = y[:, 1:].reshape(-1)\n",
    "\n",
    "            # 忽略pad\n",
    "            select = y != vocab_y['<PAD>']\n",
    "            pred = pred[select]\n",
    "            y = y[select]\n",
    "\n",
    "            loss = loss_func(pred, y)\n",
    "            optim.zero_grad()\n",
    "            loss.backward()\n",
    "            optim.step()\n",
    "\n",
    "            # 打印信息\n",
    "            if i % 2000 == 0:\n",
    "                #[select, 39] -> [select]\n",
    "                pred = pred.argmax(1)\n",
    "                correct = (pred == y).sum().item()\n",
    "                accuracy = correct / len(pred)\n",
    "                lr = optim.param_groups[0]['lr']\n",
    "                print(epoch, i, lr, loss.item(), accuracy)\n",
    "\n",
    "        sched.step()\n",
    "\n",
    "\n",
    "train()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eb9368cf",
   "metadata": {},
   "source": [
    "## 测试"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "f5898f5c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "<SOS>mbpfuom5hmz7nv7pzt5csdkbvazzkxvpngmmkyxfh5iu<EOS><PAD><PAD><PAD><PAD>\n",
      "<SOS>UUI4HFXYKMMGNPVXKZZAVBKDSC4TZP2VN2ZMH4MOUFPBM<EOS><PAD><PAD><PAD><PAD>\n",
      "<SOS>UUI4HFXYKMMGNPVXKZZAVBKDSC4TZP2VN2ZMH4MOUFPBM<EOS><EOS><EOS><EOS>\n",
      "1\n",
      "<SOS>yvkhhyj5uh6hbllz8ihkrjxxbyzucxn1ogmlahx9cdngz<EOS><PAD><PAD><PAD>\n",
      "<SOS>ZZGNDC0XHALMGO8NXCUZYBXXJRKHI1ZLLBH3HU4JYHHKVY<EOS><PAD><PAD><PAD>\n",
      "<SOS>ZZGNDC0XHALMGO8NXCUZYBXXJRKHI1ZLLBH3HU4JYHHKVY<EOS><EOS><EOS>\n",
      "2\n",
      "<SOS>pzvzkxbb7mxh8snvb76pgxpcmlmjbndxjpd<EOS><PAD><PAD><PAD><PAD><PAD><PAD><PAD><PAD><PAD><PAD><PAD><PAD><PAD>\n",
      "<SOS>DDPJXDNBJMLMCPXGP32BVNS1HXM2BBXKZVZP<EOS><PAD><PAD><PAD><PAD><PAD><PAD><PAD><PAD><PAD><PAD><PAD><PAD><PAD>\n",
      "<SOS>DDPJXDNBJMLMCPXGP32BVNS1HXM2BBXKZVZP<EOS><EOS><EOS><EOS><EOS><EOS><EOS><EOS><EOS><EOS><EOS><EOS><EOS>\n",
      "3\n",
      "<SOS>aflmsnhjxzdhjn1dvi9bmrj5nzmkmpad9dm<EOS><PAD><PAD><PAD><PAD><PAD><PAD><PAD><PAD><PAD><PAD><PAD><PAD><PAD>\n",
      "<SOS>MMD0DAPMKMZN4JRMB0IVD8NJHDZXJHNSMLFA<EOS><PAD><PAD><PAD><PAD><PAD><PAD><PAD><PAD><PAD><PAD><PAD><PAD><PAD>\n",
      "<SOS>MMD0DAPMKMZN4JRMB0IVD8NJHDZXJHNSMLFA<EOS><EOS><EOS><EOS><EOS><EOS><EOS><EOS><EOS><EOS><EOS><EOS><EOS>\n",
      "4\n",
      "<SOS>c9nnvljekhmh6vfba8bxzacv7glmlyme6x<EOS><PAD><PAD><PAD><PAD><PAD><PAD><PAD><PAD><PAD><PAD><PAD><PAD><PAD><PAD>\n",
      "<SOS>XX3EMYLMLG2VCAZXB1ABFV3HMHKEJLVNN0C<EOS><PAD><PAD><PAD><PAD><PAD><PAD><PAD><PAD><PAD><PAD><PAD><PAD><PAD><PAD>\n",
      "<SOS>XX3EMYLMLG2VCAZXB1ABFV3HMHKEJLVNN0C<EOS><EOS><EOS><EOS><EOS><EOS><EOS><EOS><EOS><EOS><EOS><EOS><EOS><EOS>\n",
      "5\n",
      "<SOS>mkhzhfkovnknmc4zlfvkbuj6lbxxjnxxbhhbk5vxzvxhj<EOS><PAD><PAD><PAD>\n",
      "<SOS>JJHXVZXV4KBHHBXXNJXXBL3JUBKVFLZ5CMNKNVOKFHZHKM<EOS><PAD><PAD><PAD>\n",
      "<SOS>JJHXVZXV4KBHHBXXNJXXBL3JUBKVFLZ5CMNKNVOKFHZHKM<EOS><EOS><EOS>\n",
      "6\n",
      "<SOS>jnacbaxmadmzxj7jxomvnhsr1xbo8pskjb5b<EOS><PAD><PAD><PAD><PAD><PAD><PAD><PAD><PAD><PAD><PAD><PAD><PAD>\n",
      "<SOS>BB4BJKSP1OBX8RSHNVMOXJ2JXZMDAMXABCANJ<EOS><PAD><PAD><PAD><PAD><PAD><PAD><PAD><PAD><PAD><PAD><PAD><PAD>\n",
      "<SOS>BB4BJKSP1OBX8RSHNVMOXJ2JXZMDAMXABCANJ<EOS><EOS><EOS><EOS><EOS><EOS><EOS><EOS><EOS><EOS><EOS><EOS>\n",
      "7\n",
      "<SOS>pcozkz7kjmoodzbyn0uzzac4zggncr5l36nl4dxah9vvfzx<EOS><PAD>\n",
      "<SOS>XXZFVV0HAXD5LN36L4RCNGGZ5CAZZU9NYBZDOOMJK2ZKZOCP<EOS><PAD>\n",
      "<SOS>XXZFVV0HAXD5LN36L4RCNGGZ5CAZZU9NYBZDOOMJK2ZKZOCP<EOS>\n"
     ]
    }
   ],
   "source": [
    "#第13章/测试\n",
    "def test():\n",
    "    for i, (x, y) in enumerate(loader):\n",
    "        break\n",
    "\n",
    "    for i in range(8):\n",
    "        print(i)\n",
    "        print(''.join([vocab_xr[i] for i in x[i].tolist()]))\n",
    "        print(''.join([vocab_yr[i] for i in y[i].tolist()]))\n",
    "        print(''.join(\n",
    "            [vocab_yr[i] for i in predict(x[i].unsqueeze(0))[0].tolist()]))\n",
    "\n",
    "\n",
    "test()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
